from torch_geometric.nn import GATConv, SAGPooling, GCNConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
from torch_geometric.nn import MessagePassing
from torch_geometric.nn.conv.gcn_conv import GCNConv
import torch.nn as nn
import torch
import torch.nn.functional as F
from typing import List, Tuple

from models.model_metrics import cal_metrics

# 本地库


class VarmisuseForRepairLayer(nn.Module):
    def __init__(
        self,
        out_features,
        task_type,
        max_node_per_graph=5,
        criterion=nn.CrossEntropyLoss(),
        metrics=cal_metrics,
        device="cpu",
    ):
        super(VarmisuseForRepairLayer, self).__init__()
        
        self.max_node_per_graph = max_node_per_graph
        
        self.varmisuse_base_linear=nn.Sequential(
            nn.Linear(out_features,out_features),
            nn.LayerNorm(out_features)
        )
        
        self.varmiuse_repair_bug=nn.Sequential(
            nn.Linear(2*out_features,1)
        )
        
        # self.varmisuse_layer = nn.Sequential(
            
        # )
        
        self.task_type=task_type
        self.criterion=criterion
        self.metrics = metrics
        self.device=device

    def run_for_repair(self,output,candidate_ids,candidate_masks,slot_id,location_index):
        location_embeding=torch.stack([output[i][location_index[i]] for i in range(output.size()[0])])

        #print(location_embeding.size())
        
        candidate_embedding=torch.einsum("bcd,bc->bcd",output,candidate_ids)
        
        #slot_inner_product = torch.einsum("bd,bcd->bc", location_embeding, candidate_embedding)  #shape g, c

        slot_embedding_unsqueeze = torch.unsqueeze(location_embeding, dim=1)  # shape: g,1,d
        slot_embedding_repeat = slot_embedding_unsqueeze.repeat(1, self.max_node_per_graph, 1)  # shape: g, c, d

        # print(slot_embedding_repeat.size())
        # print(slot_embedding_unsqueeze.size())
        # print(slot_inner_product.size())

        slot_cand_comb = torch.cat(
            [candidate_embedding, slot_embedding_repeat], dim=2)  #shape: g, c, d*2
        repair_vector = self.varmiuse_repair_bug(slot_cand_comb)  # shape: g, c, 1
        repair_vector = torch.squeeze(repair_vector, dim=-1)  # shape: g, c
        
        #print(repair_vector.size())
        
        repair_vector += (1.0 - candidate_masks.view(-1, self.max_node_per_graph)) * -1e7
        
        return repair_vector
        
    def forward(self,
                output,
                slot_id,
                candidate_ids,
                candidate_masks,
                label=None,
                output_label=False,
                metric_functions=None,
                **kwargs):
        
        output=self.varmisuse_base_linear(output)
        
        repair_predict=None
        

        repair_predict=self.run_for_repair(output,candidate_ids,candidate_masks,slot_id,label)

            
        return None,repair_predict